Skip to content

Keep SnapKV window scores on device 🤖🤖🤖 - #316

Closed
0z5a wants to merge 1 commit into
NVIDIA:mainfrom
0z5a:snapkv-device-sentinel
Closed

0z5a wants to merge 1 commit into
NVIDIA:mainfrom
0z5a:snapkv-device-sentinel

Conversation

@0z5a

@0z5a 0z5a commented Nov 4, 2026

Copy link
Copy Markdown

PR description

Keep SnapKV's observation-window sentinel on the device instead of reading scores.max().item() back to the host. This removes that host synchronization while retaining the global maximum plus one and the existing score shape.

Validation

  • Targeted score and CUDA synchronization tests: 7 passed each on B300 and two GH200 systems.
  • Compared a fixed-weight, randomly initialized small Llama before and after this change: FP16/BF16 with eager and SDPA attention produced identical prefill logits, compressed keys and values, and one decode step. This does not establish pretrained-model accuracy or end-to-end performance.
  • make style passed; the full make test suite was not run.

Checklist

  • Tests are working (make test) — targeted tests passed; full suite not run
  • Code is formatted correctly (make style)
  • Copyright header is included
  • All commits are signed off

🤖🤖🤖

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@copy-pr-bot

copy-pr-bot Bot commented Nov 4, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@SimJeg

SimJeg commented Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Hi, this does not change accuracy and is negligible for speed

@SimJeg SimJeg closed this Nov 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants